# -*- coding: utf-8 -*-
import akshare as ak
import pandas as pd
import numpy as np
import datetime
import os
import argparse
import time
import random
from functools import lru_cache

# 导入安全请求模块
from akshare_request import (
    get_stock_zh_a_hist_safe,
    get_stock_zh_a_spot_em_safe,
    get_tool_trade_date_hist_sina_safe
)

# 配置参数, 最近时间内涨停的股票排名
ANALYSIS_CONFIG = {
    'days': 10,                # 查询最近的交易日数量
    'top_n': 100,              # 获取涨停次数排名前N的股票
    'min_market_value': 1000000000,  # 最小市值（10亿）
    'max_market_value': 500000000000,  # 最大市值（5000亿）
    'include_st': False,       # 是否包含ST股票
    'include_gem': False,       # 是否包含创业板（30开头）
    'include_sci_tech': False,  # 是否包含科创板（68开头）
    'include_main': True,      # 是否包含主板（00、60开头）
    'output_folder': 'data'    # 输出文件夹
}

def format_date(date_str):
    """格式化日期字符串为YYYYMMDD格式"""
    # 如果输入的是YYYY-MM-DD格式，转换为YYYYMMDD
    if '-' in date_str:
        return date_str.replace('-', '')
    return date_str

def get_stock_list():
    """获取符合条件的A股股票列表，使用安全请求函数"""
    try:
        # 使用安全请求获取A股股票基本信息
        stock_info = get_stock_zh_a_spot_em_safe()
        
        # 根据配置筛选股票代码
        code_filters = []
        if ANALYSIS_CONFIG['include_main']:
            code_filters.extend(["00", "60"])
        if ANALYSIS_CONFIG['include_gem']:
            code_filters.append("30")
        if ANALYSIS_CONFIG['include_sci_tech']:
            code_filters.append("68")
            
        if not code_filters:  # 如果没有选择任何板块，返回空DataFrame
            print("错误：至少需要选择一个板块（主板、创业板或科创板）")
            return pd.DataFrame()
            
        # 筛选股票代码
        stock_list = stock_info[stock_info["代码"].str.startswith(tuple(code_filters))]
        
        # 筛选ST股票
        if not ANALYSIS_CONFIG['include_st']:
            stock_list = stock_list[~stock_list["名称"].str.contains("ST")]
        
        # 筛选市值范围
        stock_list["流通市值"] = pd.to_numeric(stock_list["流通市值"], errors="coerce")
        stock_list = stock_list[
            (stock_list["流通市值"] >= ANALYSIS_CONFIG['min_market_value']) & 
            (stock_list["流通市值"] <= ANALYSIS_CONFIG['max_market_value'])
        ]
        
        return stock_list
    except Exception as e:
        print(f"获取股票列表失败: {e}")
        return pd.DataFrame()

def get_trading_dates(days=10):
    """获取最近N个交易日的日期列表，使用安全请求函数"""
    try:
        # 获取当前日期
        today = datetime.datetime.now().date()  # 转换为date对象
        # 计算30天前的日期（为了确保能获取到足够的交易日）
        days_ago_30 = today - datetime.timedelta(days=30)
        
        # 使用安全请求获取交易日历
        trade_cal = get_tool_trade_date_hist_sina_safe()
        # 筛选指定日期范围内的交易日
        trade_dates = trade_cal[(trade_cal['trade_date'] >= days_ago_30) & 
                               (trade_cal['trade_date'] <= today)]
        
        # 获取最近N个交易日
        recent_dates = trade_dates.tail(days)['trade_date'].tolist()
        
        # 确保日期格式一致 (YYYYMMDD) - 转换为akshare的stock_zh_a_hist函数所需格式
        formatted_dates = []
        for date in recent_dates:
            # 将datetime.date对象转换为YYYYMMDD格式字符串
            date_str = date.strftime('%Y%m%d')
            formatted_dates.append(date_str)
        
        return formatted_dates
    except Exception as e:
        print(f"获取交易日期失败: {e}")
        return []

@lru_cache(maxsize=300)
def get_stock_history(stock_code, start_date, end_date, max_retries=3):
    """获取指定股票在给定时间范围内的历史数据，使用安全请求函数"""
    try:
        # 格式化日期
        start_date = format_date(start_date)
        end_date = format_date(end_date)
        
        # 使用安全请求函数获取历史数据
        # 该函数内部已包含重试、限流和缓存机制
        df = get_stock_zh_a_hist_safe(
            symbol=stock_code, 
            period="daily", 
            start_date=start_date, 
            end_date=end_date, 
            adjust="qfq"
        )
        
        return df if not df.empty else pd.DataFrame()
    except Exception as e:
        print(f"{stock_code}: 获取历史数据失败 {e}")
        return pd.DataFrame()

def is_limit_up(row, stock_code):
    """判断是否涨停"""
    # 根据股票代码判断涨停幅度
    if stock_code.startswith(('300', '688')):
        # 创业板、科创板涨停幅度为20%
        limit_pct = 20
    else:
        # 主板涨停幅度为10%
        limit_pct = 10
        
    # 计算涨幅 (akshare返回的数据中可能是'涨跌幅'或'涨跌幅(%)')
    if '涨跌幅' in row:
        change_pct = row['涨跌幅']
    elif '涨跌幅(%)' in row:
        change_pct = row['涨跌幅(%)']
    else:
        # 如果没有直接的涨跌幅列，尝试计算
        try:
            if '开盘' in row and '收盘' in row:
                # 使用当日收盘价与开盘价计算涨幅
                change_pct = (row['收盘'] - row['开盘']) / row['开盘'] * 100
            else:
                # 无法计算涨幅
                return False
        except:
            return False
    
    # 判断是否涨停（允许有0.1%的误差）
    return change_pct >= (limit_pct - 0.1)

def count_limit_ups(stock_code, stock_name, trading_dates):
    """统计指定股票在给定交易日内的涨停次数"""
    try:
        if not trading_dates:
            return None
            
        # 获取开始日期和结束日期
        start_date = trading_dates[0]
        end_date = trading_dates[-1]
        
        # 获取历史数据
        df = get_stock_history(stock_code, start_date, end_date)
        if df.empty:
            return None
        
        # 计算涨停次数
        df['是否涨停'] = df.apply(lambda row: is_limit_up(row, stock_code), axis=1)
        limit_up_count = df['是否涨停'].sum()
        
        # 获取涨停日期列表
        # 确定日期列名
        date_column = '日期'
        if '日期' not in df.columns and 'trade_date' in df.columns:
            date_column = 'trade_date'
            
        # 获取涨停日期并格式化
        limit_up_dates_raw = df[df['是否涨停']][date_column].tolist()
        limit_up_dates = []
        
        # 格式化日期为YYYY-MM-DD
        for date in limit_up_dates_raw:
            if isinstance(date, str):
                # 如果日期是字符串格式
                if len(date) == 8 and date.isdigit():  # YYYYMMDD格式
                    formatted_date = f"{date[:4]}-{date[4:6]}-{date[6:]}"
                    limit_up_dates.append(formatted_date)
                else:  # 其他格式，保持原样
                    limit_up_dates.append(date)
            else:  # 如果是其他类型（如datetime），转为字符串
                limit_up_dates.append(str(date))
        
        # 计算区间涨幅
        start_price = df.iloc[0]["收盘"]
        end_price = df.iloc[-1]["收盘"]
        growth_rate = (end_price - start_price) / start_price * 100
        
        # 计算平均成交量
        avg_volume = df["成交量"].mean()
        
        return {
            "代码": stock_code,
            "名称": stock_name,
            "涨停次数": int(limit_up_count),
            "涨停日期": ", ".join(limit_up_dates) if limit_up_dates else "",  # 使用逗号+空格分隔，提高可读性
            "区间涨幅(%)": round(growth_rate, 2),
            "平均成交量": int(avg_volume),
            "开始价格": round(start_price, 2),
            "结束价格": round(end_price, 2),
            "交易天数": len(df)
        }
    except Exception as e:
        print(f"{stock_code}: 统计涨停次数失败 {e}")
        return None

def analyze_limit_ups(days=10, top_n=100, batch_size=50):
    """分析最近N个交易日内涨停次数最多的股票，使用批处理减少并发请求"""
    # 获取最近N个交易日
    trading_dates = get_trading_dates(days)
    if not trading_dates:
        print("获取交易日期失败，程序结束")
        return pd.DataFrame()
        
    # 格式化显示日期（YYYY-MM-DD）
    start_display = trading_dates[0][:4] + '-' + trading_dates[0][4:6] + '-' + trading_dates[0][6:] if len(trading_dates[0]) == 8 else trading_dates[0]
    end_display = trading_dates[-1][:4] + '-' + trading_dates[-1][4:6] + '-' + trading_dates[-1][6:] if len(trading_dates[-1]) == 8 else trading_dates[-1]
    print(f"\n===== 分析最近{days}个交易日（{start_display}至{end_display}）涨停情况 =====\n")
    
    # 获取股票列表
    stock_list = get_stock_list()
    if stock_list.empty:
        print("获取股票列表失败，程序结束")
        return pd.DataFrame()
    
    print(f"共获取 {len(stock_list)} 只符合条件的股票")
    
    # 统计每只股票的涨停次数，使用批处理
    results = []
    total = len(stock_list)
    
    # 将股票列表分成批次处理
    for batch_start in range(0, total, batch_size):
        batch_end = min(batch_start + batch_size, total)
        batch = stock_list.iloc[batch_start:batch_end]
        
        print(f"\n处理批次: {batch_start+1}-{batch_end}/{total}")
        
        batch_results = []
        for _, row in batch.iterrows():
            stock_code = row["代码"]
            stock_name = row["名称"]
            
            result = count_limit_ups(stock_code, stock_name, trading_dates)
            if result and result["涨停次数"] > 0:  # 只保留有涨停的股票
                batch_results.append(result)
        
        # 添加到总结果
        results.extend(batch_results)
        
        # 批次间暂停，避免请求过于频繁
        if batch_end < total:
            pause_time = random.uniform(1.0, 3.0)
            print(f"批次处理完成，暂停 {pause_time:.1f} 秒...")
            time.sleep(pause_time)
    
    # 转换为DataFrame并按涨停次数排序
    if not results:
        print("没有符合条件的股票，程序结束")
        return pd.DataFrame()
    
    df_results = pd.DataFrame(results)
    df_results = df_results.sort_values(by=["涨停次数", "区间涨幅(%)"], ascending=[False, False])
    
    # 获取前N名
    top_stocks = df_results.head(top_n)
    
    return top_stocks

def save_results(df, days):
    """保存结果到文件"""
    if df.empty:
        return
    
    # 创建输出文件夹
    folder_path = os.path.join(os.path.dirname(__file__), ANALYSIS_CONFIG['output_folder'])
    os.makedirs(folder_path, exist_ok=True)
    
    # 获取当前日期
    today = datetime.datetime.now().strftime('%Y%m%d')
    
    # 生成文件名
    file_name = f"涨停排名_{days}天_{today}.txt"
    file_path = os.path.join(folder_path, file_name)
    
    # 保存到文件
    df.to_csv(file_path, sep="\t", index=False, encoding="utf-8")
    print(f"\n数据已保存至 {file_path}")

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="分析最近N个交易日内涨停次数最多的股票")
    parser.add_argument("-d", "--days", type=int, help="查询最近的交易日数量")
    parser.add_argument("-n", "--top", type=int, help="获取涨停次数排名前N的股票")
    parser.add_argument("--min-value", type=float, help="最小市值（亿元）")
    parser.add_argument("--max-value", type=float, help="最大市值（亿元）")
    parser.add_argument("--include-st", action="store_true", help="是否包含ST股票")
    parser.add_argument("--exclude-gem", action="store_true", help="是否排除创业板")
    parser.add_argument("--exclude-sci-tech", action="store_true", help="是否排除科创板")
    parser.add_argument("--exclude-main", action="store_true", help="是否排除主板")
    parser.add_argument("-o", "--output", help="输出文件夹")
    
    args = parser.parse_args()
    
    # 更新配置
    if args.days:
        ANALYSIS_CONFIG['days'] = args.days
    if args.top:
        ANALYSIS_CONFIG['top_n'] = args.top
    if args.min_value:
        ANALYSIS_CONFIG['min_market_value'] = args.min_value * 100000000  # 亿元转换为元
    if args.max_value:
        ANALYSIS_CONFIG['max_market_value'] = args.max_value * 100000000  # 亿元转换为元
    if args.include_st:
        ANALYSIS_CONFIG['include_st'] = True
    if args.exclude_gem:
        ANALYSIS_CONFIG['include_gem'] = False
    if args.exclude_sci_tech:
        ANALYSIS_CONFIG['include_sci_tech'] = False
    if args.exclude_main:
        ANALYSIS_CONFIG['include_main'] = False
    if args.output:
        ANALYSIS_CONFIG['output_folder'] = args.output

def main():
    """主函数"""
    # 解析命令行参数
    parse_arguments()
    
    start_time = time.time()
    print(f"开始分析: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"分析最近 {ANALYSIS_CONFIG['days']} 个交易日的涨停情况")
    print(f"筛选条件: 市值范围 {ANALYSIS_CONFIG['min_market_value']/100000000:.1f}亿 - {ANALYSIS_CONFIG['max_market_value']/100000000:.1f}亿")
    print(f"包含ST股票: {'是' if ANALYSIS_CONFIG['include_st'] else '否'}")
    print(f"包含板块: {'主板 ' if ANALYSIS_CONFIG['include_main'] else ''}{'创业板 ' if ANALYSIS_CONFIG['include_gem'] else ''}{'科创板' if ANALYSIS_CONFIG['include_sci_tech'] else ''}")
    
    try:
        # 分析股票
        top_stocks = analyze_limit_ups(
            ANALYSIS_CONFIG['days'],
            ANALYSIS_CONFIG['top_n']
        )
        
        # 显示结果
        if not top_stocks.empty:
            print(f"\n===== 涨停次数排名前 {ANALYSIS_CONFIG['top_n']} 的股票 =====\n")
            print(top_stocks.to_string(index=False))
            
            # 保存结果
            save_results(top_stocks, ANALYSIS_CONFIG['days'])
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
    finally:
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"\n分析完成，耗时: {elapsed_time:.2f} 秒 ({elapsed_time/60:.2f} 分钟)")
        print(f"结束时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")


if __name__ == "__main__":
    main()